import cv2
import keyboard
import threading
import sys, os
import time
import pygame
import json
import shutil
from pathlib import Path
import socket
import base64

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))


class control:
    def __init__(self, speed):
        self.img = []
        self.speed = speed
        self.quit = False
        self.y_button = 0  # 0是没输入,1是右,-1是左
        self.y_state = 'up'
        self.angle_button = 0  # 0是没输入,1是右,-1是左
        self.angle_state = 'up'
        self.y_dict = {"a": -1, "d": 1}
        self.angle_dict = {"left": -1, "right": 1}
        self.y_speed = 0
        self.angle_speed = 0

    def start(self):
        # 钩住所有的按键事件
        b = threading.Thread(target=self.thread)
        b.setDaemon(True)
        b.start()
        img_catch = threading.Thread(target=self.catch)
        img_catch.setDaemon(True)
        img_catch.start()

    def thread(self):
        keyboard.hook(self.det_button)
        keyboard.wait()

    def det_button(self, event):
        if event.name in self.y_dict:
            self.update_y_button(event)
        elif event.name == "w":
            if event.event_type == "down":
                self.speed = 0.1
            else:
                self.speed = 0
        elif event.name in self.angle_dict:
            self.update_angle_button(event)
        elif event.name in "q":
            self.quit = True
        # time.sleep(0.1)

    def update_y_button(self, event):
        if self.y_state == "up" or self.y_button == self.y_dict[event.name]:
            self.y_state = event.event_type
            if self.y_state == "down":
                self.y_button = self.y_dict[event.name]
            else:
                self.y_button = 0
                self.y_speed = 0

    def update_angle_button(self, event):
        if self.angle_state == "up" or self.angle_button == self.angle_dict[event.name]:
            self.angle_state = event.event_type
            if self.angle_state == "down":
                self.angle_button = self.angle_dict[event.name]
            else:
                self.angle_button = 0
                self.angle_speed = 0

    def catch(self):
        while True and not self.quit:
            cap = cv2.VideoCapture("/dev/video-front")
            ret, frame = cap.read()
            if ret:
                self.img.append([frame, self.speed, self.y_speed, self.angle_speed])


def restore(controller, dir):
    json_text = []
    n = 1
    while True:
        if not os.path.exists(dir + str(n)):
            os.makedirs(dir + str(n))
            dir = dir + str(n)
            break
        else:
            n += 1

    for i, data_list in enumerate(controller.img):
        img_name = f"/{i}.jpg"
        x_speed = data_list[1]
        y_speed = data_list[2]
        angle_speed = data_list[3]
        state = [x_speed, y_speed, angle_speed]
        json_text.append({
            "img_path": img_name[1:],
            "state": state
        })
        img = data_list[0]
        cv2.imwrite(dir + img_name, img)
    with open(dir + "/data_label.json", 'w') as f:
        json.dump(json_text, f)
    print(f"{dir}保存成功")


def play_controller1():
    from vehicle import MecanumBase
    car = MecanumBase()
    y_speed_index = 0.000001
    angle_speed_index = 0.00001
    controller = control(0)
    controller.start()
    while True:
        if controller.quit:
            break
        if controller.y_state == "down":
            controller.y_speed += y_speed_index * controller.y_button
        if controller.angle_state == "down":
            controller.angle_speed += angle_speed_index * controller.angle_button
        # sys.stdout.write(f"\ry button {controller.y_button} is {controller.y_state},angle button {controller.angle_button} is {controller.angle_state}")
        sys.stdout.write(f"\r [{controller.speed},{controller.y_speed},{controller.angle_speed}]")
        sys.stdout.flush()
        car.mecanum_wheel(controller.speed, controller.y_speed, controller.angle_speed)
    car.mecanum_wheel(0, 0, 0)
    while True:
        decision = input("是否保存这轮数据(y/n):")
        if decision == "y":
            restore(controller, "img_dataset/")
        elif decision == "n":
            break
        else:
            print("输错了!")
            continue


class control2:
    def __init__(self, car, speed, fps=4):
        self.dir = ''
        self.fps = fps
        self.car = car
        self.json_text = []
        self.img = 0
        self.collect_speed = speed
        self.speed = 0
        self.quit = False
        self.state = "control"
        self.y_speed = 0
        self.angle_speed = 0
        self.end = False
        self.save = 0
        # self.client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        # self.client.connect(("192.168.116.86", 5050))

    def start(self):
        # 钩住所有的按键事件
        b = threading.Thread(target=self.thread)
        b.setDaemon(True)
        b.start()
        img_catch = threading.Thread(target=self.catch)
        img_catch.setDaemon(True)
        img_catch.start()

    def thread(self):
        # 初始化pygame
        pygame.init()
        # 初始化操纵杆模块
        pygame.joystick.init()
        # 检查操纵杆数量
        joystick_count = pygame.joystick.get_count()
        print(f"Number of joysticks: {joystick_count}")
        # 如果有操纵杆，创建一个操纵杆对象
        if joystick_count > 0:
            joystick = pygame.joystick.Joystick(0)
            joystick.init()  # 初始化操纵杆
            print(f"Joystick name: {joystick.get_name()}")
        while not self.quit:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    self.quit = True
                elif event.type == pygame.JOYBUTTONDOWN:
                    if event.button == 3:  # 按Y退出全部程序
                        self.quit = True
                    if event.button == 1 and self.state == "collect":  # 按B不保存数据切换状态
                        self.speed = 0
                        self.state = "control"
                        time.sleep(0.1)
                        self.clear()
                        print("已切换成控制模式")
                    if event.button == 0 and self.state == "collect":  # 按A保存数据并切换状态
                        self.speed = 0
                        self.state = "control"
                        self.store()
                        print("已切换成控制模式")
                    if event.button == 2 and self.state == "control":  # 按X进入采集模式
                        self.mkdir()
                        self.state = "collect"
                        self.speed = self.collect_speed
                        print("已切换成采集模式")
                elif event.type == pygame.JOYAXISMOTION:
                    if event.axis == 0:
                        self.y_speed = event.value * 0.1
                    if event.axis == 1 and self.state == "control":
                        self.speed = event.value * (-1) * self.collect_speed
                    if event.axis == 3:
                        self.angle_speed = event.value

    def store(self):
        with open(self.dir + "/data_label.json", 'w') as f:
            json.dump(self.json_text, f)
        print(f"{self.dir} 保存成功")

    def mkdir(self):
        dir = 'data_collect/dataset'
        n = 1
        while True:
            if not os.path.exists(dir + str(n)):
                os.makedirs(dir + str(n))
                self.dir = dir + str(n)
                break
            else:
                n += 1

    def clear(self):
        # 删除目录及其内容
        directory = Path(self.dir)

        # 确保目录存在
        if directory.exists() and directory.is_dir():
            for item in directory.iterdir():
                if item.is_dir():
                    shutil.rmtree(item)
                else:
                    item.unlink()
            directory.rmdir()

    def send(self, img):
        _, buffer = cv2.imencode('.jpg', img)
        jpg_as_text = base64.b64encode(buffer).decode('utf-8')
        length = str(len(jpg_as_text)).encode('utf-8')
        length += b' ' * (16 - len(length))
        message = length + jpg_as_text.encode('utf-8')
        self.client.sendall(message)

    def catch(self):
        cap = cv2.VideoCapture("/dev/video-front")
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG'))
        cap.set(cv2.CAP_PROP_FPS, self.fps)  # 控制采集数量,正常30帧一圈大概3000张
        while True and not self.quit:
            # ret, frame = cap.read()
            # if ret:
            #     self.send(frame)
            if self.state == "control":
                self.img = 0
                self.json_text = []
                continue
            elif self.state == "collect":
                ret, frame = cap.read()
                if ret:
                    img_name = f"/{self.img}.jpg"
                    state = [self.speed, self.y_speed, self.angle_speed]
                    self.json_text.append({"img_path": img_name[1:], "state": state})
                    cv2.imwrite(self.dir + img_name, frame)
                    self.img += 1
                time.sleep(0.25)


class control3:
    def __init__(self, car, speed):
        self.dir = ''
        self.car = car
        self.json_text = []
        self.img = 0
        self.collect_speed = speed
        self.speed = 0
        self.quit = False
        self.state = "control"
        self.y_speed = 0
        self.angle_speed = 0
        self.end = False
        self.save = 0
        # self.client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        # self.client.connect(("192.168.116.86", 5050))

    def start(self):
        # 钩住所有的按键事件
        b = threading.Thread(target=self.thread)
        b.setDaemon(True)
        b.start()
        img_catch = threading.Thread(target=self.catch)
        img_catch.setDaemon(True)
        img_catch.start()

    def thread(self):
        # 初始化pygame
        pygame.init()
        # 初始化操纵杆模块
        pygame.joystick.init()
        # 检查操纵杆数量
        joystick_count = pygame.joystick.get_count()
        print(f"Number of joysticks: {joystick_count}")
        # 如果有操纵杆，创建一个操纵杆对象
        if joystick_count > 0:
            joystick = pygame.joystick.Joystick(0)
            joystick.init()  # 初始化操纵杆
            print(f"Joystick name: {joystick.get_name()}")
        while not self.quit:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    self.quit = True
                elif event.type == pygame.JOYBUTTONDOWN:
                    if event.button == 3:  # 按Y退出全部程序
                        self.quit = True
                    if event.button == 0 and self.state == "collect":  # 按B不保存数据切换状态
                        self.speed = 0
                        self.state = "control"
                        time.sleep(0.1)
                        self.clear()
                        print("已切换成控制模式")
                    if event.button == 1 and self.state == "collect":  # 按A保存数据并切换状态
                        self.speed = 0
                        self.state = "control"
                        self.store()
                        print("已切换成控制模式")
                    if event.button == 4 and self.state == "control":  # 按X进入采集模式
                        self.mkdir()
                        self.state = "collect"
                        self.speed = self.collect_speed
                        print("已切换成采集模式")
                elif event.type == pygame.JOYAXISMOTION:
                    if event.axis == 0:
                        self.y_speed = event.value * 0.1
                    if event.axis == 1 and self.state == "control":
                        self.speed = event.value * (-0.15)
                    if event.axis == 2:
                        self.angle_speed = event.value

    def store(self):
        with open(self.dir + "/data_label.json", 'w') as f:
            json.dump(self.json_text, f)
        print(f"{self.dir} 保存成功")

    def mkdir(self):
        dir = 'data_collect/dataset'
        n = 1
        while True:
            if not os.path.exists(dir + str(n)):
                os.makedirs(dir + str(n))
                self.dir = dir + str(n)
                break
            else:
                n += 1

    def clear(self):
        # 删除目录及其内容
        directory = Path(self.dir)

        # 确保目录存在
        if directory.exists() and directory.is_dir():
            for item in directory.iterdir():
                if item.is_dir():
                    shutil.rmtree(item)
                else:
                    item.unlink()
            directory.rmdir()

    def send(self, img):
        _, buffer = cv2.imencode('.jpg', img)
        jpg_as_text = base64.b64encode(buffer).decode('utf-8')
        length = str(len(jpg_as_text)).encode('utf-8')
        length += b' ' * (16 - len(length))
        message = length + jpg_as_text.encode('utf-8')
        self.client.sendall(message)

    def catch(self):
        cap = cv2.VideoCapture("/dev/video-front")
        cap.set(cv2.CAP_PROP_FPS, 4)  # 控制采集数量,正常30帧一圈大概3000张
        while True and not self.quit:
            # ret, frame = cap.read()
            # if ret:
            #     self.send(frame)
            if self.state == "control":
                self.img = 0
                self.json_text = []
                continue
            elif self.state == "collect":
                ret, frame = cap.read()
                if ret:
                    img_name = f"/{self.img}.jpg"
                    state = [self.speed, self.y_speed, self.angle_speed]
                    self.json_text.append({"img_path": img_name[1:], "state": state})
                    cv2.imwrite(self.dir + img_name, frame)
                    self.img += 1


def play_controller2():
    from vehicle import MecanumBase
    car = MecanumBase()
    controller = control2(car, 0.15, 30)
    controller.start()
    print("按Y退出全部程序,按B保存数据并切换模式,按A不保存数据并切换模式")
    time.sleep(0.5)
    while True:
        if controller.quit:
            break
        sys.stdout.write(
            f"\r img:{controller.img} - [{controller.speed},{controller.y_speed},{controller.angle_speed}]                                              ")
        sys.stdout.flush()
        car.mecanum_wheel(controller.speed, controller.y_speed, controller.angle_speed)


def play_controller3():
    from vehicle import MecanumBase
    car = MecanumBase()
    controller = control3(car, 0.15)
    controller.start()
    print("按Y退出全部程序,按B保存数据并切换模式,按A不保存数据并切换模式")
    time.sleep(0.5)
    while True:
        if controller.quit:
            break
        sys.stdout.write(
            f"\r img:{controller.img} - [{controller.speed},{controller.y_speed},{controller.angle_speed}]                                              ")
        sys.stdout.flush()
        car.mecanum_wheel(controller.speed, controller.y_speed, controller.angle_speed)


def collect_tasks(dir, pages):
    cap = cv2.VideoCapture(0)
    for i in range(1, 14):
        n = 100
        print(f"现在采集第{i}个任务的数据")
        while True:
            if n == pages:  # 张数采集足够就退出,进入下一个任务
                break
            ret, frame = cap.read()  # 读取摄像头
            if ret:
                cv2.imshow("img", frame)  # 展示图片
                if cv2.waitKey(1) & 0xFF == ord('s'):  # 按下s保存
                    cv2.imwrite(f"{dir}/{i}_{n}.jpg", frame)
                    n += 1
                    print(f"任务{i}保存了第{n}张图片,在{dir}/{i}_{n}.jpg")
        cont = input("下一个任务?")


if __name__ == "__main__":
    # play_controller1()  # 键盘
    play_controller2()  # 手柄采集
    # play_controller3()  # 连线
    # collect_tasks("task_det_data",200)
